use crate::byte::{BitReader, ByteRead};

use super::{FseResult, FseDistribution};

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct FseCell {
    pub base_line: u32,
    pub num_bits: u8,
    pub symbol: u8,
}

#[derive(Debug, PartialEq)]
pub struct FseTable {
    table:Vec<FseCell>,
    accuracy_log:u8,
}

impl FseTable {
    pub fn first_cell<'a,R:ByteRead>(&self,reader:&mut BitReader<'a,R>) -> FseResult<FseCell> {
        let state = reader.read_reverse_bits::<usize>(self.accuracy_log as usize)?;
        Ok(self.table[state])
    }

    pub fn next_cell<'a,R:ByteRead>(&self,reader:&mut BitReader<'a,R>,entry:FseCell) -> FseResult<FseCell>{
        let addition = reader.read_reverse_bits::<u8>(entry.num_bits as usize)?;
        let new_state = entry.base_line + addition as u32;
        Ok(self.table[new_state as usize])
    }
}

impl std::fmt::Display for FseTable {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let table_string = self.table
            .iter()
            .enumerate()
            .map(|(i, cell)|format!("{}\t{}\t{}\t\t{}",i,cell.symbol,cell.base_line,cell.num_bits))
            .collect::<Vec<_>>()
            .join("\n");
        write!(f,"accuracy log:{}\nstate\tsymbol\tbase_line\tnum_bits\n{}", self.accuracy_log, table_string)
    }
}

impl From<FseDistribution> for FseTable {
    fn from(value: FseDistribution) -> Self {
        let empty_cell = FseCell { base_line: 0, num_bits: 0, symbol: 0 };
        let table_size = 1 << value.accuracy_log;
        let mut symbol_counter = vec![0;value.probabilities.len()];
        let mut decoding_table = vec![empty_cell;table_size];
        let mut less_than_one_position = decoding_table.len();
        let mut natural_position = 0;
        value.probabilities
            .iter()
            .enumerate()
            .filter(|(_,&prob)|prob.eq(&-1))
            .for_each(|(i,_)|{
                less_than_one_position -= 1;
                decoding_table[less_than_one_position] = FseCell {
                    base_line: 0, num_bits: value.accuracy_log, symbol: i as u8
                };
            });
        value.probabilities
            .iter()
            .enumerate()
            .filter(|(_,&prob)|!prob.eq(&-1))
            .for_each(|(i,&prob)|{
                (0..prob).into_iter().for_each(|_|{
                    decoding_table[natural_position].symbol = i as u8;
                    natural_position = next_cell_position_skip(natural_position, table_size, less_than_one_position);
                });
            });
        decoding_table
            .iter_mut()
            .enumerate()
            .filter(|(i, _)|i.lt(&less_than_one_position))
            .for_each(|(_, cell)|{
                let counter = &mut symbol_counter[cell.symbol as usize];
                let probabilities = value.probabilities[cell.symbol as usize];
                let (base_line, num_bits) = cal_base_line_and_num_bits(table_size, probabilities, *counter);
                cell.base_line = base_line;
                cell.num_bits = num_bits;
                *counter += 1;
            });
        Self { table:decoding_table, accuracy_log:value.accuracy_log }
    }
}

impl From<&[(usize,u8,u8,u32)]> for FseTable {
    fn from(value: &[(usize,u8,u8,u32)]) -> Self {
        let accuracy_log = 8 * std::mem::size_of::<usize>() as u8 - 1 - value.len().leading_zeros() as u8;
        let table = value
            .iter()
            .map(|(_, symbol, num_bits, base_line)|
                FseCell{ symbol:*symbol, base_line:*base_line, num_bits:*num_bits })
            .collect();
        Self { table, accuracy_log }
    }
}

fn next_cell_position_skip(position:usize, table_size:usize, bound:usize) -> usize {
    let mut position = position;
    loop {
        position = next_cell_position(position, table_size);
        if position < bound { break position; }
    }
}

fn next_cell_position(position:usize,table_size:usize) -> usize {
    let position = position + (table_size >> 1) + (table_size >> 3) + 3;
    let position = position & (table_size - 1);
    position
}

fn ceil_2_power(num:u32) -> u32 {
    assert!(num!=0);
    if (num & (num - 1))!=0 {let mut val = num;
        val |= val>>1;
        val |= val>>2;
        val |= val>>4;
        val |= val>>8;
        val |= val>>16;
        val + 1
    }else{
        num
    }
}

fn cal_base_line_and_num_bits(table_size:usize, probabilities:i32, order:u32) -> (u32,u8) {
    let table_size = table_size as u32;
    let probabilities = probabilities as u32;
    let part_num = ceil_2_power(probabilities as u32);
    let share = table_size / part_num;
    let double_share_num = (table_size - share * probabilities)/share;
    if order < double_share_num {
        let base_line = (probabilities - double_share_num) * share + 2 * order * share;
        let num_bits = 8 * std::mem::size_of_val(&share) as u8 - share.leading_zeros() as u8;
        (base_line, num_bits)
    } else {
        let base_line = (order - double_share_num) * share;
        let num_bits = 8 * std::mem::size_of_val(&share) as u8 - 1 - share.leading_zeros() as u8;
        (base_line, num_bits)
    }
}

pub fn test_predefined_fse_decoding_table() {
    let default_ll_distribution = FseDistribution::predefined_literals_length_distribution();
    let default_ml_distribution = FseDistribution::predefined_match_length_distribution();
    let default_oc_distribution = FseDistribution::predefined_offset_codes_distribution();
    let ll_table:FseTable = default_ll_distribution.into();
    let ml_table:FseTable = default_ml_distribution.into();
    let oc_table:FseTable = default_oc_distribution.into();
    let real_ll_table:FseTable = (&super::PREDEFINED_LL_TABLE[..]).into();
    let real_ml_table:FseTable = (&super::PREDEFINED_ML_TABLE[..]).into();
    let real_oc_table:FseTable = (&super::PREDEFINED_OC_TABLE[..]).into();
    //println!("literals length:{}",ll_table);
    assert_eq!(ll_table,real_ll_table);
    //println!("match length:{}",ml_table);
    assert_eq!(ml_table,real_ml_table);
    //println!("offset codes:{}",oc_table);
    assert_eq!(oc_table,real_oc_table);
}